# check_dataset.py
# usage: 带有 sql 的json文件，执行脚本后，屏幕显示结果集
# python ./src/tools/check_dataset.py --src "./result/part1_results.jsonl" --id 2001

import argparse
import json
import sqlite3
import os
from typing import List, Tuple

def parse_args():
    parser = argparse.ArgumentParser(description="根据ID从源文件提取SQL并在本地数据库中执行")
    parser.add_argument('--src', type=str, required=True, help='源JSONL文件路径')
    parser.add_argument('--id', type=int, required=True, help='要查找的记录ID')
    return parser.parse_args()

def find_sql_by_id(src_file: str, target_id: int) -> str:
    """在 src 文件中查找指定 id 的 sql 语句"""
    if not os.path.exists(src_file):
        raise FileNotFoundError(f"源文件未找到: {src_file}")

    with open(src_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
                if record.get('id') == target_id:
                    # 打印 question 字段（如果存在）
                    if 'question' in record:
                        print(f"问题: {record['question']}")
                    return record.get('sql')
            except json.JSONDecodeError as e:
                print(f"解析JSON失败: {line}")
                continue
#                raise e
    raise ValueError(f"未找到 id={target_id} 的记录")

def execute_sql_on_db(sql: str, db_path: str) -> Tuple[List[str], List[Tuple]]:
    """连接数据库并执行SQL，返回 (列名列表, 数据列表)"""
    if not os.path.exists(db_path):
        raise FileNotFoundError(f"数据库文件未找到: {db_path}")

    conn = None
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        cursor.execute(sql)
        
        # 获取列名
        columns = [desc[0] for desc in cursor.description] if cursor.description else []
        
        # 获取所有数据
        rows = cursor.fetchall()
        return columns, rows
    except sqlite3.Error as e:
        raise RuntimeError(f"SQL执行出错: {e}\nSQL: {sql}")
    finally:
        if conn:
            conn.close()

def main():
    args = parse_args()
    
    # 定义数据库路径
    db_path = './data/data.db'
    
    try:
        # 步骤1: 从 src 文件中查找对应 id 的 SQL
        sql = find_sql_by_id(args.src, args.id)
        if not sql:
            print(f"ID为 {args.id} 的记录没有SQL语句。")
            return
        
        print(f"✅ 找到 ID={args.id} 的SQL:")
        print(f"   {sql}\n")
        
        # 步骤2: 执行SQL
        columns, rows = execute_sql_on_db(sql, db_path)
        
        # 步骤3: 输出结果
        if not columns:
            print("⚠️  查询无结果（无字段）")
            return
        
        print("📊 查询结果:")
        print(" | ".join(f"{col:^15}" for col in columns))
        print("-" * (len(columns) * 18))
        
        if not rows:
            print(" (无数据)")
        else:
            for row in rows:
                print(" | ".join(f"{str(val):^15}" for val in row))
                
    except Exception as e:
        print(f"❌ 错误: {e}")
        exit(1)

if __name__ == '__main__':
    main()